#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
宇宙平衡控制器
确保系统在各种状态下的动态平衡
"""

import paddle
import paddle.nn as nn
import paddle.nn.functional as F
import numpy as np
from typing import Dict, Tuple, Optional

class UniversalBalanceController(nn.Layer):
    """
    宇宙平衡控制器
    维护系统的阴阳动态平衡
    """
    
    def __init__(self, 
                 input_dim: int = 512,
                 output_dim: int = 512):
        """
        初始化平衡控制器
        
        Args:
            input_dim: 输入维度
            output_dim: 输出维度
        """
        super().__init__()
        
        self.input_dim = input_dim
        self.output_dim = output_dim
        
        # 平衡检测网络
        self.balance_detector = BalanceDetector(input_dim)
        
        # 动态调节网络
        self.dynamic_adjuster = DynamicAdjustmentNetwork(input_dim, output_dim)
        
        # 稳定性保证器
        self.stability_ensurer = StabilityEnsurer(output_dim)
        
        # 平衡历史记录
        self.balance_history = []
    
    def forward(self, x: paddle.Tensor) -> paddle.Tensor:
        """
        平衡控制流程
        
        Args:
            x: 输入状态
            
        Returns:
            平衡后的状态
        """
        # 检测当前平衡状态
        balance_analysis = self.balance_detector(x)
        
        # 动态调节
        adjusted = self.dynamic_adjuster(x, balance_analysis)
        
        # 稳定性保证
        stabilized = self.stability_ensurer(adjusted)
        
        # 记录平衡状态
        self.record_balance_state(balance_analysis, stabilized)
        
        return stabilized
    
    def record_balance_state(self, analysis: Dict, output: paddle.Tensor):
        """记录平衡状态"""
        record = {
            'balance_score': analysis['balance_score'],
            'stability_level': analysis['stability_level'],
            'adjustment_magnitude': analysis['adjustment_needed'],
            'output_norm': paddle.norm(output).item(),
            'output_entropy': self.calculate_entropy(output)
        }
        self.balance_history.append(record)
    
    def calculate_entropy(self, state: paddle.Tensor) -> float:
        """计算状态熵值"""
        prob_dist = F.softmax(state, axis=-1)
        entropy = -paddle.sum(prob_dist * paddle.log(prob_dist + 1e-8))
        return entropy.item()
    
    def get_balance_report(self) -> Dict:
        """获取平衡报告"""
        if not self.balance_history:
            return {}
        
        latest = self.balance_history[-1]
        history_scores = [record['balance_score'] for record in self.balance_history]
        
        avg_balance = np.mean(history_scores)
        trend = self.analyze_balance_trend(history_scores)
        
        return {
            'current_balance': latest['balance_score'],
            'average_balance': avg_balance,
            'balance_trend': trend,
            'stability': latest['stability_level'],
            'adjustment_history': len(self.balance_history),
            'recent_records': self.balance_history[-10:]  # 最近10条记录
        }
    
    def analyze_balance_trend(self, scores: List[float]) -> str:
        """分析平衡趋势"""
        if len(scores) < 2:
            return "稳定"
        
        recent_scores = scores[-5:]  # 最近5个分数
        if len(recent_scores) < 2:
            return "稳定"
        
        # 简单趋势分析
        differences = [recent_scores[i] - recent_scores[i-1] 
                      for i in range(1, len(recent_scores))]
        avg_difference = np.mean(differences)
        
        if avg_difference > 0.01:
            return "改善中"
        elif avg_difference < -0.01:
            return "下降中"
        else:
            return "稳定"

class BalanceDetector(nn.Layer):
    """平衡检测器"""
    
    def __init__(self, input_dim: int):
        super().__init__()
        self.input_dim = input_dim
        
        # 平衡特征提取
        self.feature_extractor = nn.Sequential(
            nn.Linear(input_dim, input_dim * 2),
            nn.GELU(),
            nn.Linear(input_dim * 2, input_dim // 2)
        )
        
        # 平衡评分网络
        self.balance_scorer = nn.Sequential(
            nn.Linear(input_dim // 2, 64),
            nn.ReLU(),
            nn.Linear(64, 1),
            nn.Sigmoid()
        )
        
        # 稳定性评估
        self.stability_evaluator = nn.Linear(input_dim // 2, 3)  # 3个稳定性级别
    
    def forward(self, x: paddle.Tensor) -> Dict:
        """检测平衡状态"""
        # 提取平衡特征
        features = self.feature_extractor(x)
        
        # 计算平衡分数
        balance_score = self.balance_scorer(features).squeeze(-1)
        
        # 评估稳定性
        stability_logits = self.stability_evaluator(features)
        stability_probs = F.softmax(stability_logits, axis=-1)
        stability_level = paddle.argmax(stability_probs, axis=-1).item()
        
        # 确定需要的调节程度
        adjustment_needed = 1.0 - balance_score.item()
        
        stability_labels = {0: "不稳定", 1: "基本稳定", 2: "高度稳定"}
        
        return {
            'balance_score': balance_score.item(),
            'stability_level': stability_labels.get(stability_level, "未知"),
            'stability_confidence': paddle.max(stability_probs).item(),
            'adjustment_needed': adjustment_needed,
            'balance_features': features
        }

class DynamicAdjustmentNetwork(nn.Layer):
    """动态调节网络"""
    
    def __init__(self, input_dim: int, output_dim: int):
        super().__init__()
        self.input_dim = input_dim
        self.output_dim = output_dim
        
        # 阳调节网络（增强创造性）
        self.yang_adjuster = nn.Sequential(
            nn.Linear(input_dim, output_dim * 2),
            nn.Tanh(),
            nn.Linear(output_dim * 2, output_dim)
        )
        
        # 阴调节网络（增强稳定性）
        self.yin_adjuster = nn.Sequential(
            nn.Linear(input_dim, output_dim * 2),
            nn.Sigmoid(), 
            nn.Linear(output_dim * 2, output_dim)
        )
        
        # 动态权重网络
        self.weight_controller = nn.Linear(input_dim, 2)
    
    def forward(self, x: paddle.Tensor, balance_analysis: Dict) -> paddle.Tensor:
        """动态调节"""
        adjustment_needed = balance_analysis['adjustment_needed']
        
        # 计算阴阳调节权重
        weights = F.softmax(self.weight_controller(x), axis=-1)
        yang_weight, yin_weight = weights[0], weights[1]
        
        # 应用调节
        yang_adjusted = self.yang_adjuster(x)
        yin_adjusted = self.yin_adjuster(x)
        
        # 基于平衡需求混合
        base_mix = yang_weight * yang_adjusted + yin_weight * yin_adjusted
        
        # 根据调节需求调整
        if adjustment_needed > 0.5:
            # 需要大幅调节，增加动态性
            dynamic_factor = adjustment_needed
            final_output = dynamic_factor * base_mix + (1 - dynamic_factor) * x
        else:
            # 小幅调节，保持稳定性
            final_output = 0.3 * base_mix + 0.7 * x
        
        return final_output

class StabilityEnsurer(nn.Layer):
    """稳定性保证器"""
    
    def __init__(self, dim: int):
        super().__init__()
        self.dim = dim
        
        # 稳定性检测
        self.stability_check = nn.Linear(dim, 1)
        
        # 稳定性增强
        self.stability_enhancer = nn.Sequential(
            nn.Linear(dim, dim * 2),
            nn.GELU(),
            nn.Linear(dim * 2, dim)
        )
    
    def forward(self, x: paddle.Tensor) -> paddle.Tensor:
        """保证输出稳定性"""
        # 检测稳定性
        stability_score = F.sigmoid(self.stability_check(x))
        
        if stability_score.item() < 0.7:
            # 需要稳定性增强
            enhanced = self.stability_enhancer(x)
            # 混合原始和增强版本
            mix_ratio = stability_score.item()
            stabilized = mix_ratio * enhanced + (1 - mix_ratio) * x
        else:
            # 已经稳定，直接返回
            stabilized = x
        
        # 最终归一化
        return F.normalize(stabilized)

class AdaptiveBalanceController(nn.Layer):
    """
    自适应平衡控制器
    根据系统状态自动调整平衡策略
    """
    
    def __init__(self, input_dim: int, output_dim: int):
        super().__init__()
        self.input_dim = input_dim
        self.output_dim = output_dim
        
        # 多个平衡策略
        self.balance_strategies = nn.LayerList([
            UniversalBalanceController(input_dim, output_dim),
            UniversalBalanceController(input_dim, output_dim),
            UniversalBalanceController(input_dim, output_dim)
        ])
        
        # 策略选择器
        self.strategy_selector = nn.Linear(input_dim, len(self.balance_strategies))
        
        # 策略融合器
        self.strategy_fusion = nn.Linear(output_dim * len(self.balance_strategies), output_dim)
    
    def forward(self, x: paddle.Tensor) -> paddle.Tensor:
        """自适应平衡控制"""
        # 选择最佳策略
        strategy_weights = F.softmax(self.strategy_selector(x), axis=-1)
        
        # 应用所有策略
        strategy_outputs = []
        for i, strategy in enumerate(self.balance_strategies):
            output = strategy(x)
            strategy_outputs.append(output)
        
        # 加权融合
        weighted_outputs = []
        for i, output in enumerate(strategy_outputs):
            weight = strategy_weights[:, i].unsqueeze(-1)
            weighted_outputs.append(output * weight)
        
        # 合并输出
        combined = paddle.concat(weighted_outputs, axis=-1)
        final_output = self.strategy_fusion(combined)
        
        return F.normalize(final_output)